package lox

import java.util.*

class Resolver(
    private val interpreter: Interpreter,
    private val scopes: Stack<MutableMap<String, Boolean>> = Stack(),
    private var currentFunction: FunctionType = FunctionType.NONE,
    private var currentClass: ClassType = ClassType.NONE,
) : Expr.Visitor<Unit>, Stmt.Visitor<Unit> {

    enum class FunctionType {
        NONE, FUNCTION, INITIALIZER, METHOD
    }

    enum class ClassType {
        NONE, CLASS, SUBCLASS,
    }

    override fun visitBlockStmt(stmt: Stmt.Block) {
        beginScope()
        resolve(stmt.statements)
        endScope()
    }

    override fun visitClassStmt(stmt: Stmt.Class) {

        val enclosingClass: ClassType = ClassType.CLASS
        currentClass = ClassType.CLASS

        declare(stmt.name)
        define(stmt.name)

        if (stmt.superclass != null && stmt.superclass.name.lexeme == stmt.name.lexeme) {
            Lox.error(stmt.superclass.name, "A class can't inherit from itself.")
        }

        stmt.superclass?.apply {
            currentClass = ClassType.SUBCLASS
            resolve(this)
        }

        if (stmt.superclass != null) {
            beginScope()
            scopes.peek()["super"] = true
        }

        beginScope()
        scopes.peek()["this"] = true
        for (method in stmt.methods) {
            val declaration = if (method.name.lexeme == "init") FunctionType.INITIALIZER else FunctionType.METHOD
            resolveFunction(method, declaration)
        }
        endScope()

        if (stmt.superclass != null) endScope()

        currentClass = enclosingClass
    }

    override fun visitExpressionStmt(stmt: Stmt.Expression) = resolve(stmt.expression)

    override fun visitIfStmt(stmt: Stmt.If) {
        resolve(stmt.condition)
        resolve(stmt.thenBranch)
        if (stmt.elseBranch != null) resolve(stmt.elseBranch)
    }

    override fun visitPrintStmt(stmt: Stmt.Print) = resolve(stmt.expression)

    override fun visitReturnStmt(stmt: Stmt.Return) {
        if (currentFunction == FunctionType.NONE) Lox.error(stmt.keyword, "Can't return from top-level code.")
        if (stmt.value != null) {
            if (currentFunction == FunctionType.INITIALIZER) Lox.error(
                stmt.keyword, "Can't return a value from an initializer."
            )

            resolve(stmt.value)
        }
    }

    override fun visitWhileStmt(stmt: Stmt.While) {
        resolve(stmt.condition)
        resolve(stmt.body)
    }

    override fun visitFunctionStmt(stmt: Stmt.Function) {
        declare(stmt.name)
        define(stmt.name)

        resolveFunction(stmt, FunctionType.FUNCTION)
    }

    private fun resolveFunction(function: Stmt.Function, type: FunctionType) {

        val enclosingFunction = currentFunction
        currentFunction = type

        beginScope()
        function.params.forEach { param -> declare(param); define(param) }
        resolve(function.body)
        endScope()

        currentFunction = enclosingFunction
    }

    override fun visitVarStmt(stmt: Stmt.Var) {
        declare(stmt.name)
        stmt.initializer?.let { resolve(it) }
        define(stmt.name)
    }

    private fun define(name: Token) {
        if (!scopes.isEmpty()) scopes.peek()[name.lexeme] = true
    }

    private fun declare(name: Token) {
        if (!scopes.isEmpty()) scopes.peek().also {
            if (it.containsKey(name.lexeme)) Lox.error(
                name, "Already variable with this name in this scope"
            )
        }[name.lexeme] = false
    }

    private fun endScope() {
        scopes.pop()
    }

    private fun beginScope() {
        scopes.push(mutableMapOf())
    }

    fun resolve(statements: List<Stmt>) = statements.forEach(::resolve)

    private fun resolve(statement: Stmt) = statement.accept(this)
    private fun resolve(expr: Expr) = expr.accept(this)
    override fun visitAssignExpr(expr: Expr.Assign) {
        resolve(expr.value)
        resolveLocal(expr, expr.name)
    }

    override fun visitBinaryExpr(expr: Expr.Binary) {
        resolve(expr.left)
        resolve(expr.right)
    }

    override fun visitCallExpr(expr: Expr.Call) {
        resolve(expr.callee)
        expr.arguments.forEach(::resolve)
    }

    override fun visitGetExpr(expr: Expr.Get) = resolve(expr.obj)

    override fun visitGroupingExpr(expr: Expr.Grouping) = resolve(expr.expression)

    override fun visitLiteralExpr(expr: Expr.Literal) = Unit

    override fun visitLogicalExpr(expr: Expr.Logical) {
        resolve(expr.left)
        resolve(expr.right)
    }

    override fun visitSetExpr(expr: Expr.Set) {
        resolve(expr.value)
        resolve(expr.obj)
    }

    override fun visitSuperExpr(expr: Expr.Super) = when (currentClass) {
        ClassType.NONE -> Lox.error(expr.keyword, "Can't use 'super' out of a class.")
        ClassType.CLASS -> Lox.error(expr.keyword, "Can't use 'super' in a class with no superclass.")
        else -> resolveLocal(expr, expr.keyword)
    }

    override fun visitThisExpr(expr: Expr.This) = when (currentClass) {
        ClassType.NONE -> Lox.error(expr.keyword, "Can't use 'this' out of a class.")
        ClassType.CLASS, ClassType.SUBCLASS -> resolveLocal(expr, expr.keyword)
    }

    override fun visitUnaryExpr(expr: Expr.Unary) = resolve(expr.right)

    override fun visitVariableExpr(expr: Expr.Variable) {
        if (!scopes.isEmpty() && scopes.peek()[expr.name.lexeme] == false) Lox.error(
            expr.name, "Can't read local variable in its own initializer."
        )

        resolveLocal(expr, expr.name)
    }

    private fun resolveLocal(expr: Expr, name: Token) {
        for ((index, scope) in scopes.reversed().withIndex()) {
            if (scope.containsKey(name.lexeme)) {
                interpreter.resolve(expr, index)
                return
            }
        }
    }

}